1951B - Battle Cows - CodeForces Solution

def calculate(n,k,A):
    track = -1
    cow_rating = A[k-1]
    res = res2 = 0
 
    for i in range(k-1):
        if A[i]>cow_rating:
            track = i
            break
        
    temp = A[0]
    A[0] = A[k-1]
    A[k-1] = temp
 
    for i in range(1,n):
        if A[i]<cow_rating:
            res+=1
        elif A[i]>cow_rating:
            break
 
    temp = A[0]
    A[0] = A[k-1]
    A[k-1] = temp
    
    if track!=-1 and track!=0:
        temp = A[k-1]
        A[k-1] = A[track]
        A[track] = A[k-1]
 
        res2 += 1
        for i in range(track+1, n):
            if A[i]<cow_rating:
                res2+=1
            elif A[i]>cow_rating:
                break
 
    return max(res, res2)
 
for i in range(int(input())):
    n,k = list(map(int, input().split()))
    A = list(map(int, input().split()))
    print(calculate(n,k,A))